package com.gateway.filter;

import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.RequestPath;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;
import java.util.Map;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setResponseStatus;

/**
 * 限流过滤器
 * 目前报异常【ERR Error running script (call to f_9d491aea731237273f4274f9ed9660b432b23791): @user_script:48: ERR value is not an integer or out of range 】
 * @author dzh
 * @date 2023/2/1 14:12
 */
@Slf4j
@Component
public class RateLimiterGlobalFilter implements GlobalFilter {

    @Resource
    private RedisRateLimiter redisRateLimiter;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        RequestPath requestPath = exchange.getRequest().getPath();
        String path = requestPath.toString();
        Route route = exchange
                .getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
        String routeId = route.getId();
        return redisRateLimiter.isAllowed(routeId, path).flatMap(response -> {

            for (Map.Entry<String, String> header : response.getHeaders()
                    .entrySet()) {
                exchange.getResponse().getHeaders().add(header.getKey(),
                        header.getValue());
            }

            if (response.isAllowed()) {
                return chain.filter(exchange);
            }

            log.info("rateLimiter not allowed {}", routeId);
            setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
            return exchange.getResponse().setComplete();
        });
    }
}
